#include "kernel_loader.h"
#include <iostream>
#include <vector>
#include <cassert>
#include <cmath>
#include <sys/time.h>
#include <acl/acl.h>
#include <experiment/runtime/runtime/rt.h>
#include "kernel_launchers.h"  // 包含所有注册的 kernel launchers
// 测试参数
const int NUMEL = 98432;
#define DTYPE float

// 时间函数
static unsigned long long currentTime() {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (tv.tv_sec * 1000000 + tv.tv_usec);
}

void initDeviceMemory(void* ptr, size_t size, float value) {
    std::vector<float> hostData(size / sizeof(float), value);
    rtError_t error = rtMemcpy(ptr, size, hostData.data(), size, RT_MEMCPY_HOST_TO_DEVICE);
    assert(error == RT_ERROR_NONE);
}

bool verifyResult(void* devPtr, size_t size, float expectedValue) {
    std::vector<float> hostData(size / sizeof(float));
    rtError_t error = rtMemcpy(hostData.data(), size, devPtr, size, RT_MEMCPY_DEVICE_TO_HOST);
    assert(error == RT_ERROR_NONE);
    
    for (size_t i = 0; i < hostData.size(); ++i) {
        if (std::abs(hostData[i] - expectedValue) > 1e-6) {
            std::cout << "Mismatch at index " << i << ": expected " << expectedValue 
                      << ", got " << hostData[i] << std::endl;
            return false;
        }
    }
    return true;
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cout << "Usage: " << argv[0] << " <kernel_binary_path> <kernel_name>" << std::endl;
        std::cout << "Example: " << argv[0] << " ./kernels/add_kernel.o add_kernel" << std::endl;
        return -1;
    }

    std::string binaryPath = argv[1];
    std::string kernelName = argv[2];

    std::cout << "###### Start Add Kernel Test ######" << std::endl;
    std::cout << "Binary: " << binaryPath << std::endl;
    std::cout << "Kernel: " << kernelName << std::endl;

    // 1. 初始化ACL和设备
    aclInit(nullptr);
    int32_t device_id = 0;
    rtError_t error = rtSetDevice(device_id);
    if (error != RT_ERROR_NONE) {
        std::cout << "Failed to set device: error=" << error << std::endl;
        return -1;
    }
    std::cout << "Device initialized: " << device_id << std::endl;

    // 2. 加载kernel
    auto kernelHandle = KernelLoader::loadKernel(kernelName, binaryPath);
    if (!kernelHandle.isValid()) {
        std::cout << "Failed to load kernel: " << kernelName << std::endl;
        return -1;
    }
    std::cout << "Kernel loaded: " << kernelHandle.getName() << std::endl;

    // 3. 创建执行流
    rtStream_t stream;
    error = rtStreamCreate(&stream, 0);
    assert(error == RT_ERROR_NONE);
    std::cout << "Stream created" << std::endl;

    // 4. 分配设备内存
    void* input1 = nullptr;
    void* input2 = nullptr; 
    void* output = nullptr;
    
    error = rtMalloc(&input1, NUMEL * sizeof(DTYPE), RT_MEMORY_HBM, 2);
    assert(error == RT_ERROR_NONE);
    error = rtMalloc(&input2, NUMEL * sizeof(DTYPE), RT_MEMORY_HBM, 2);
    assert(error == RT_ERROR_NONE);
    error = rtMalloc(&output, NUMEL * sizeof(DTYPE), RT_MEMORY_HBM, 2);
    assert(error == RT_ERROR_NONE);
    std::cout << "Device memory allocated" << std::endl;

    // 5. 初始化输入数据
    initDeviceMemory(input1, NUMEL * sizeof(DTYPE), 2.0f);  // input1 = 2.0
    initDeviceMemory(input2, NUMEL * sizeof(DTYPE), 3.0f);  // input2 = 3.0
    std::cout << "Input data initialized: input1=2.0, input2=3.0" << std::endl;

    // 6. 准备 kernel 参数
    void* workspace_addr = nullptr;  // 传入 nullptr，launcher 会根据 JSON 配置自动分配
    void* syncBlockLock = nullptr;   // 传入 nullptr，launcher 会根据 JSON 配置自动分配
    
    int32_t gridX = 97, gridY = 1, gridZ = 1;
    std::cout << "Kernel arguments prepared" << std::endl;
    std::cout << "  Grid config: " << gridX << " x " << gridY << " x " << gridZ << std::endl;

    // 7. 执行kernel（使用生成的函数接口）
    std::cout << "Launching kernel..." << std::endl;
    unsigned long long startTime = currentTime();
    
    error = TritonLauncher::add_kernel(
        stream,
        gridX, gridY, gridZ,
        workspace_addr,
        syncBlockLock,
        input1,   // arg0
        input2,   // arg1
        output,   // arg2
        NUMEL     // arg3
    );
    
    if (error != RT_ERROR_NONE) {
        std::cout << "add_kernel failed: error=" << error << std::endl;
        return -1;
    }

    // 9. 同步等待完成
    error = rtStreamSynchronize(stream);
    if (error != RT_ERROR_NONE) {
        std::cout << "Stream synchronize failed: error=" << error << std::endl;
        return -1;
    }

    unsigned long long endTime = currentTime();
    std::cout << "Kernel execution completed" << std::endl;
    std::cout << "  Time cost: " << (endTime - startTime) / 1000.0 << " ms" << std::endl;

    // 10. 验证结果
    std::cout << "Verifying results..." << std::endl;
    bool success = verifyResult(output, NUMEL * sizeof(DTYPE), 5.0f);  // 2.0 + 3.0 = 5.0
    
    if (success) {
        std::cout << "Test PASSED: 2.0 + 3.0 = 5.0" << std::endl;
    } else {
        std::cout << "Test FAILED: Result verification failed" << std::endl;
        // 打印部分结果用于调试
        std::vector<float> hostData(NUMEL);
        rtMemcpy(hostData.data(), NUMEL * sizeof(float), output, NUMEL * sizeof(float), RT_MEMCPY_DEVICE_TO_HOST);
        std::cout << "First 10 results: ";
        for (int i = 0; i < 10 && i < NUMEL; ++i) {
            std::cout << hostData[i] << " ";
        }
        std::cout << std::endl;
    }

    // 11. 清理资源
    std::cout << "Cleaning up resources..." << std::endl;
    rtFree(input1);
    rtFree(input2);
    rtFree(output);
    rtStreamDestroy(stream);
    
    KernelLoader::cleanup();
    rtDeviceReset(device_id);
    aclFinalize();

    std::cout << "###### Test Completed ######" << std::endl;
    return success ? 0 : -1;
}